package reverse_proxy

import (
	"context"
	"io"
	"log"
	"net"
	"time"

	"gin-gateway/reverse_proxy/load_balance"
)

// 自定义的 tcp reverseproxy
func NewTcpLoadBalanceReverseProxy(c *TcpSliceRouterContext, lb load_balance.LoadBalancer) *TcpReverseProxy {
	return func() *TcpReverseProxy {
		nextAddr, err := lb.Get("")
		if err != nil || nextAddr == "" {
			panic("get next addr fail")
		}
		return &TcpReverseProxy{
			ctx:             c.Ctx,
			Addr:            nextAddr,
			KeepAlivePeriod: time.Second,
			DialTimeout:     time.Second,
		}
	}()
}

// TcpReverseProxy TCP反向代理
type TcpReverseProxy struct {
	ctx             context.Context //单次请求单独设置
	Addr            string
	KeepAlivePeriod time.Duration //设置
	DialTimeout     time.Duration //设置超时时间
	OnDialError     func(src net.Conn, dstDialErr error)
}

func (tp *TcpReverseProxy) dialTimeout() time.Duration {
	if tp.DialTimeout > 0 {
		return tp.DialTimeout
	}
	return 5 * time.Second
}

func (tp *TcpReverseProxy) keepAlivePeriod() time.Duration {
	if tp.KeepAlivePeriod != 0 {
		return tp.KeepAlivePeriod
	}
	return time.Minute
}

// 传入上游 conn，在这里建立下游连接，
// 用 io.Copy 在两者之间做数据交换
func (tp *TcpReverseProxy) ServeTCP(ctx context.Context, src net.Conn) {
	// 与下游建立连接
	dst, err := (&net.Dialer{
		Timeout:   tp.dialTimeout(),     //连接超时
		KeepAlive: tp.keepAlivePeriod(), //设置连接的检测时长
	}).DialContext(ctx, "tcp", tp.Addr)

	if err != nil {
		log.Printf("TcpReverseProxy: for incoming conn %v, error dialing %q: %v", src.RemoteAddr().String(), tp.Addr, err)
		tp.onDialError()(src, err)
		// 通过 panic 向外层传递错误
		panic(err)
	}
	// 关闭下游连接
	defer func() { go dst.Close() }()

	//下游连接的操作系统保活消息发送设置
	if ka := tp.keepAlivePeriod(); ka > 0 {
		if c, ok := dst.(*net.TCPConn); ok {
			err = c.SetKeepAlive(true)
			if err != nil {
				tp.onDialError()(src, err)
				panic(err)
			}
			err = c.SetKeepAlivePeriod(ka)
			if err != nil {
				tp.onDialError()(src, err)
			}
		}
	}

	_, err = src.Write([]byte("ok\n"))
	if err != nil {
		tp.onDialError()(src, err)
		panic(err)
	}

	// 在上游下游之间交换数据
	errc := make(chan error, 1)
	go tp.proxyCopy(errc, src, dst)
	go tp.proxyCopy(errc, dst, src)
	<-errc
}

func (tp *TcpReverseProxy) proxyCopy(errc chan<- error, dst, src net.Conn) {
	_, err := io.Copy(dst, src)
	errc <- err
}

func (tp *TcpReverseProxy) onDialError() func(src net.Conn, dstDialErr error) {
	if tp.OnDialError != nil {
		return tp.OnDialError
	}
	return func(src net.Conn, dstDialErr error) {
		// log.Printf("tcpproxy: for incoming conn %v, error dialing %q: %v", src.RemoteAddr().String(), dp.Addr, dstDialErr)
		src.Close()
	}
}
